#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


import tempfile
import numpy
import scipy.linalg
import h5py
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.outcore import _load_from_h5g
from pyscf.df.addons import make_auxmol
from pyscf import __config__

MAX_MEMORY = getattr(__config__, 'df_outcore_max_memory', 2000)  # 2GB
LINEAR_DEP_THR = getattr(__config__, 'df_df_DF_lindep', 1e-12)

def sym_cderi(mol, auxbasis, int2c):
    ''' construct symmetrized (P|Q)^(-1/2) in RI'''

    auxmol = make_auxmol(mol, auxbasis)
    auxmol.build()
    symm_auxorb = auxmol.symm_orb
    irrep_id = numpy.array(auxmol.irrep_id) % 10
    nirrep = 8
    aux_nsym = numpy.zeros(nirrep, dtype = int)

    vs = []
    for irrep in range(nirrep):
        if numpy.count_nonzero(irrep_id==irrep) == 0:
            continue
        symaux = numpy.hstack([symm_auxorb[i] for i in numpy.where(irrep_id==irrep)[0]])

        cdij = symaux.T.dot(int2c).dot(symaux)
        w, v = scipy.linalg.eigh(cdij)
        idx = w > 1.e-12
        v = (v[:,idx] / numpy.sqrt(w[idx]))
        v = symaux.dot(v)
        vs.append(v)
        naux_sym = v.shape[1]
        aux_nsym[irrep] = naux_sym
#
    vsym = numpy.hstack(vs)
    return aux_nsym, vsym
    

def cholesky_eri_b(mol, erifile, auxbasis='weigend+etb', dataname='j3c',
                   int3c='int3c2e', aosym='s2ij', int2c='int2c2e', comp=1,
                   max_memory=MAX_MEMORY, auxmol=None, verbose=logger.NOTE):
    '''3-center 2-electron DF tensor. Similar to cholesky_eri while this
    function stores DF tensor in blocks.
    '''
    assert(aosym in ('s1', 's2ij'))
    log = logger.new_logger(mol, verbose)
    time0 = (logger.process_clock(), logger.perf_counter())

    if auxmol is None:
        auxmol = make_auxmol(mol, auxbasis)
    j2c = auxmol.intor(int2c, hermi=1)
    log.debug('size of aux basis %d', j2c.shape[0])
    time1 = log.timer('2c2e', *time0)
#   try:
#       low = scipy.linalg.cholesky(j2c, lower=True)
#       tag = 'cd'
#   except scipy.linalg.LinAlgError:
#wf w, v = scipy.linalg.eigh(j2c)
#wf idx = w > LINEAR_DEP_THR
#wf low = (v[:,idx] / numpy.sqrt(w[idx]))
#wf v = None
#wf tag = 'eig'
#
    tag = 'eig'
    aux_nsym, low = sym_cderi(mol, auxbasis, j2c)

    j2c = None
    naoaux, naux = low.shape
    time1 = log.timer('Cholesky 2c2e', *time1)

    int3c = gto.moleintor.ascint3(mol._add_suffix(int3c))
    atm, bas, env = gto.mole.conc_env(mol._atm, mol._bas, mol._env,
                                      auxmol._atm, auxmol._bas, auxmol._env)
    ao_loc = gto.moleintor.make_loc(bas, int3c)
    nao = ao_loc[mol.nbas]
    naoaux = ao_loc[-1] - nao
    if aosym == 's1':
        nao_pair = nao * nao
        buflen = min(max(int(max_memory*.24e6/8/naoaux/comp), 1), nao_pair)
        shranges = _guess_shell_ranges(mol, buflen, 's1')
    else:
        nao_pair = nao * (nao+1) // 2
        buflen = min(max(int(max_memory*.24e6/8/naoaux/comp), 1), nao_pair)
        shranges = _guess_shell_ranges(mol, buflen, 's2ij')
    log.debug('erifile %.8g MB, IO buf size %.8g MB',
              naoaux*nao_pair*8/1e6, comp*buflen*naoaux*8/1e6)
    log.debug1('shranges = %s', shranges)
    # TODO: Libcint-3.14 and newer version support to compute int3c2e without
    # the opt for the 3rd index.
    #if '3c2e' in int3c:
    #    cintopt = gto.moleintor.make_cintopt(atm, mol._bas, env, int3c)
    #else:
    #    cintopt = gto.moleintor.make_cintopt(atm, bas, env, int3c)
    cintopt = gto.moleintor.make_cintopt(atm, bas, env, int3c)
    bufs1 = numpy.empty((comp*max([x[2] for x in shranges]),naoaux))

    def transform(b):
        if b.ndim == 3 and b.flags.f_contiguous:
            b = lib.transpose(b.T, axes=(0,2,1)).reshape(naoaux,-1)
        else:
            b = b.reshape((-1,naoaux)).T
        if tag == 'cd':
            if b.flags.c_contiguous:
                b = lib.transpose(b).T
            return scipy.linalg.solve_triangular(low, b, lower=True,
                                                 overwrite_b=True, check_finite=False)
        else:
            print('shape of b and low:', b.shape, low.shape)
            return lib.dot(low.T, b)

    def process(sh_range):
        bstart, bend, nrow = sh_range
        shls_slice = (bstart, bend, 0, mol.nbas, mol.nbas, mol.nbas+auxmol.nbas)
        ints = gto.moleintor.getints3c(int3c, atm, bas, env, shls_slice, comp,
                                       aosym, ao_loc, cintopt, out=bufs1)
        if comp == 1:
            dat = transform(ints)
        else:
            dat = [transform(x) for x in ints]
        return dat

    feri = _create_h5file(erifile, dataname)
    for istep, dat in enumerate(lib.map_with_prefetch(process, shranges)):
        sh_range = shranges[istep]
        label = '%s/%d'%(dataname,istep)
#       print('label', label, dat.shape)
        if comp == 1:
            feri[label] = dat
        else:
            shape = (len(dat),) + dat[0].shape
            fdat = feri.create_dataset(label, shape, dat[0].dtype.char)
            for i, b in enumerate(dat):
                fdat[i] = b
        dat = None
        log.debug('int3c2e [%d/%d], AO [%d:%d], nrow = %d',
                  istep+1, len(shranges), *sh_range)
        time1 = log.timer('gen CD eri [%d/%d]' % (istep+1,len(shranges)), *time1)
    bufs1 = None
    feri.close()
    return aux_nsym, erifile

def _guess_shell_ranges(mol, buflen, aosym):
    from pyscf.ao2mo.outcore import balance_partition
    ao_loc = mol.ao_loc_nr()
    if 's2' in aosym:
        return balance_partition(ao_loc*(ao_loc+1)//2, buflen)
    else:
        nao = ao_loc[-1]
        return balance_partition(ao_loc*nao, buflen)

def _create_h5file(erifile, dataname):
    if h5py.is_hdf5(erifile):
        feri = h5py.File(erifile, 'a')
        if dataname in feri:
            del(feri[dataname])
    else:
        feri = h5py.File(erifile, 'w')
    return feri

del(MAX_MEMORY)

